import numpy as np
import matplotlib.pyplot as plt
import matplotlib.path as mpath
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import netCDF4 as nc

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

long_lat = np.load('/home/ma-user/work/pro_data/OSI_ASTE_statistic/long_lat_file.npy')
# data = np.load('/mnt/nvme1n1/ECMWF-IFS-HR-0126-unique-grid/train/1950/sic_sit_temp_05.npy')
# data_next = np.load('/mnt/nvme1n1/ECMWF-IFS-HR-0126-unique-grid/train/1950/sic_sit_temp_06.npy')

lon = long_lat[:, 0]
lat = long_lat[:, 1]

def get_circle():
    theta = np.linspace(0, 2*np.pi, 100)
    center, radius = [0.5, 0.5], 0.5
    verts = np.vstack([np.sin(theta), np.cos(theta)]).T
    circle = mpath.Path(verts * radius + center)
    return circle


# 北极极地投影
# 创建1行2列的1个子图，使用北极地投影
FONT_DICT = {'family': "serif"}
TITLE_SIZE = 20
circle = get_circle()


def plt_metrics_sub_fig(ax, circle, data, lon=lon, lat=lat, title="Comparison"):
    ax.add_feature(cfeature.COASTLINE)  # 添加海岸线
    ax.gridlines(draw_labels=True, x_inline=False, y_inline=True, linestyle='--')  # 添加网格栅格线
    ax.set_extent([-180, 180, 60, 90], ccrs.PlateCarree())
    ax.set_title(title, fontsize=TITLE_SIZE, fontdict=FONT_DICT)
    ax.set_boundary(circle, transform=ax.transAxes)
    ax.add_feature(cfeature.COASTLINE.with_scale('110m'), lw=0.001)
    ax_s = ax.scatter(lon, lat, c=data, s=0.1, transform=ccrs.PlateCarree(), marker=".") #, edgecolors='black'
    plt.colorbar(ax_s, ax=ax, fraction=0.04)

def plt_metrics(label, pred, fig_name="comparison.png"):
    fig_window = (1, 3)
    fig = plt.figure(figsize=[30, 11])
    ax1 = fig.add_subplot(*fig_window, 1, projection=ccrs.NorthPolarStereo())
    plt_metrics_sub_fig(ax1, circle, label, title="Label")

    ax2 = fig.add_subplot(*fig_window, 2, projection=ccrs.NorthPolarStereo())
    plt_metrics_sub_fig(ax2, circle, pred, lon=lon, lat=lat, title='Prediction')

    error = label - pred
    ax3 = fig.add_subplot(*fig_window, 3, projection=ccrs.NorthPolarStereo())
    plt_metrics_sub_fig(ax3, circle, error, lon=lon, lat=lat, title="Error")

    plt.savefig(fig_name, dpi=300, bbox_inches="tight")